//
// Created by Tommy on 15/3/23.
// Copyright (c) 2015 com.koala. All rights reserved.
//

#import "NSNotificationCenter+KLObserver.h"
#import "NSObject+KLSwizzle.h"
#import "NSObject+KLExecuteOnDealloc.h"
#import <objc/runtime.h>



static char KLNSNotificationCenterArrayKey;
static char KLNSNotificationCenterAllowMethodForwardingKey;

@interface __KLObserversNotificationObserverInfo : NSObject
@property(nonatomic, copy) NSString *name;
@property(nonatomic, weak) id object;
@property(nonatomic, assign) void *blockKey;
@end

@implementation __KLObserversNotificationObserverInfo


@end

@implementation NSNotificationCenter (KLObserver)

#if KLENABLE_EXCEPTION_SWIZZLE
+ (void)load
{

    static dispatch_once_t onceToken;
    dispatch_once(&onceToken, ^{
        [self kl_swizzleInstanceSelector:@selector(addObserver:selector:name:object:) withNewSelector:KLHookedSel(addObserver:selector:name:object:)];
        [self kl_swizzleInstanceSelector:@selector(removeObserver:) withNewSelector:KLHookedSel(removeObserver:)];
        [self kl_swizzleInstanceSelector:@selector(removeObserver:name:object:) withNewSelector:KLHookedSel(removeObserver:name:object:)];
    });
}

- (BOOL)allowMethodForwarding
{
    NSNumber *state = objc_getAssociatedObject(self, &KLNSNotificationCenterAllowMethodForwardingKey);
    return [state boolValue];
}

- (void)setAllowMethodForwarding:(BOOL)allowForwarding
{
    objc_setAssociatedObject(self, &KLNSNotificationCenterAllowMethodForwardingKey,@(allowForwarding), OBJC_ASSOCIATION_RETAIN);
}

- (void)KLHookedMethod(addObserver:(id)observer selector:(SEL)aSelector name:(NSString *)aName object:(id)anObject) {
    NSMutableDictionary *registeredNotifications = (NSMutableDictionary *)objc_getAssociatedObject(observer,  &KLNSNotificationCenterArrayKey);
    if (!registeredNotifications) {
        registeredNotifications = [NSMutableDictionary dictionary];
        objc_setAssociatedObject(observer,  &KLNSNotificationCenterArrayKey, registeredNotifications, OBJC_ASSOCIATION_RETAIN);
    }

    NSMutableArray *observerInfos = registeredNotifications[NSStringFromSelector(aSelector)];
    if (!observerInfos) {
        observerInfos = [NSMutableArray array];
        registeredNotifications[NSStringFromSelector(aSelector)] = observerInfos;
    }
    __block __KLObserversNotificationObserverInfo *observerInfo = nil;

    observerInfo = [[__KLObserversNotificationObserverInfo alloc] init];
    [observerInfos addObject:observerInfo];


    observerInfo.name = aName;
    observerInfo.object = anObject;

    //! Add auto remove when observer is going to be deallocated

    __weak typeof(self) weakSelf = self;
    __weak __block id weakObject = anObject;

    void *key = [observer kl_performBlockOnDealloc:^(id obj){
        __strong typeof(weakSelf) self = weakSelf;
        id strongObserver = obj;
        NSInteger numberOfRemovals = 0;
        if ((numberOfRemovals = [self kl_removeObserver:strongObserver name:aName object:weakObject registeredNotifications:registeredNotifications])) {
            for (NSInteger i = 0; i < numberOfRemovals; ++i) {
                [self setAllowMethodForwarding:YES];
                [self KLHookedMethod(removeObserver:strongObserver  name:aName object:weakObject)];
                [self setAllowMethodForwarding:NO];
            }
        }
    }];

    //! remember the block key
    observerInfo.blockKey = key;

    //! call originalMethod
    [self KLHookedMethod(addObserver:observer selector:aSelector name:aName object:anObject)];
}


- (void)KLHookedMethod(removeObserver:(id)observer)
{
    if ([self allowMethodForwarding]) {

        [self KLHookedMethod(removeObserver:observer)];
        return;
    }

    NSMutableDictionary *registeredNotifications = (NSMutableDictionary *)objc_getAssociatedObject(observer, &KLNSNotificationCenterArrayKey);
    NSInteger numberOfRemovals = 0;
    if ((numberOfRemovals = [self kl_removeObserver:observer name:nil object:nil registeredNotifications:registeredNotifications])) {
        for (NSInteger i = 0; i < numberOfRemovals; ++i) {

            [self setAllowMethodForwarding:YES];
            [self KLHookedMethod(removeObserver:observer)];
            [self setAllowMethodForwarding:NO];
        }
    }

}

- (void)KLHookedMethod(removeObserver:(id)observer name:(NSString *)aName object:(id)anObject)
{
    if ([self allowMethodForwarding]) {
        [self KLHookedMethod(removeObserver:observer name:aName object:anObject)];
        return;
    }

    NSMutableDictionary *registeredNotifications = (NSMutableDictionary *)objc_getAssociatedObject(observer, &KLNSNotificationCenterArrayKey);
    NSUInteger numberOfRemovals = 0;
    if ([self allowMethodForwarding] || (numberOfRemovals = [self kl_removeObserver:observer name:aName object:anObject registeredNotifications:registeredNotifications])) {
        [self setAllowMethodForwarding:YES];


        [self KLHookedMethod(removeObserver:observer name:aName object:anObject)];
        [self setAllowMethodForwarding:NO];
    }

}

- (NSUInteger)kl_removeObserver:(id)observer name:(NSString *)aName object:(id)anObject registeredNotifications:(NSMutableDictionary *)registeredNotifications
{
    __block NSUInteger result = 0;

    if (aName == nil && anObject == nil) {
        //! don't need to execute block on dealloc so cleanup
        [registeredNotifications enumerateKeysAndObjectsUsingBlock:^void(id key, id obj, BOOL *stop) {
            NSMutableArray *observerInfos = obj;
            [observerInfos enumerateObjectsUsingBlock:^void(id innerObj, NSUInteger idx, BOOL *innerStop) {
                __KLObserversNotificationObserverInfo *info = innerObj;
                [observer kl_cancelDeallocBlockWithKey:info.blockKey];
            }];
        }];
        [registeredNotifications removeAllObjects];

        return 1;
    } else {
        [registeredNotifications enumerateKeysAndObjectsUsingBlock:^void(id key, id obj, BOOL *stop) {
            NSMutableArray *observerInfos = obj;
            NSMutableArray *objectsToRemove = [NSMutableArray array];
            [observerInfos enumerateObjectsUsingBlock:^void(id innerObj, NSUInteger idx, BOOL *innerStop) {
                __KLObserversNotificationObserverInfo *info = innerObj;

                if ((!aName || [aName isEqualToString:info.name]) && (!anObject || (anObject == info.object))) {
                    //! remove this info
                    [objectsToRemove addObject:innerObj];
                    //! cancel dealloc blocks
                    [observer kl_cancelDeallocBlockWithKey:info.blockKey];
                }
            }];

            //! remove all collected objects
            if ([objectsToRemove count] > 0) {

                result = [objectsToRemove count];
                [observerInfos removeObjectsInArray:objectsToRemove];
            }
        }];
    }

    return result;
}

#endif


@end
